# tests/test_columns.py
from core.extractor import LineageExtractor

def test_select_columns():
    sql = "SELECT a, b, sum(c) FROM t1;"
    st = LineageExtractor().extract(sql).statements[0]
    cols = st.columns_by_table.get("t1") or st.columns_by_table.get("*") or set()
    assert {"a", "b", "c"} <= cols
